#!/usr/bin/env python
# -*- coding=utf-8 -*-
"""make a job

Usage:
    make_job <configFile> [--outfile=<outfile>] [--logdir=<logdir>] [--queue=<queue>]

Options:
    <configFile>             The configure file
    --logdir=<logdir>        The log directory[default: ./log]
    --queue=<queue>          The queue was used[default: disease.q,disease1.q,disease2.q]
    --outfile=<outfile>      The output filename [default: out.job]

ConfigFile Example:
    #script memory  [depend_scripts] [thread]
    test1.sh 2G
    test2.sh 3G 4
    test3.sh 1G test1.sh,test2.sh
    test4.sh 1M test3.sh 5

Contact:
    suqingdong <suqingdong@novogene.com>
"""
import os
import string
import getpass
import commands
import docopt


jobTemplate = string.Template('''\
job_begin
  name ${jobname}
  memory ${memory}
  status ${status}
  sched_options ${scheds}
  cmd_begin
    ${cmd}
  cmd_end
job_end
\n''')


def write_job(jobname, memory, cmd, status='waiting', scheds='-V -cwd'):

    return jobTemplate.safe_substitute(
        jobname=jobname,
        memory=memory,
        cmd=cmd,
        status=status,
        scheds=scheds
    )


def write_order(job1, job2):

    return 'order {} after {}\n'.format(job1, job2)


def main():

    arguments = docopt.docopt(__doc__)
    # print arguments

    configFile = arguments.get('<configFile>')
    logdir = os.path.abspath(arguments.get('--logdir'))
    queue = arguments.get('--queue')
    outfile = arguments.get('--outfile')


    user = getpass.getuser()
    avaliable_queues = set(map(
        lambda x: x.split('@')[0],
        commands.getoutput('qselect -U {}'.format(user)).split('\n')
    ))

    print(avaliable_queues)

    queue_list = []
    for q in queue.split(','):
        if q in avaliable_queues:
            queue_list.append(' -q {}'.format(q))

    if not queue_list:
        print('no available queue!')
        exit()

    base_scheds = '-V -cwd' + ''.join(queue_list)

    pwd = os.getcwd()

    if not os.path.exists(configFile):
        exit('File not exists: %s' % configFile)

    order_list = []
    with open(configFile) as f, open(outfile, 'w') as out:

        for line in f:
            if line.startswith('#') or not line.strip():
                continue
            linelist = line.strip().split()
            script = linelist[0]
            memory = linelist[1]
            if script.endswith('.sh'):
                cmd = 'sh ' + os.path.join(pwd, script)
            elif script.endswith('.py'):
                cmd = 'python ' + os.path.join(pwd, script)

            thread = None
            depend_scripts = None

            if len(linelist) == 3:
                if linelist[2].isdigit():
                    thread = linelist[2]
                else:
                    depend_scripts = linelist[2]

            if len(linelist) == 4:
                depend_scripts, thread = linelist[2:4]

            # write job
            if thread:
                # cpu = int(int(thread) * 0.65)
                cpu = thread
                scheds = base_scheds + ' -l p={}'.format(cpu)
            else:
                scheds = base_scheds

            jobname = os.path.basename(script)
            out.write(write_job(jobname, memory, cmd, scheds=scheds))

            # write order
            if depend_scripts:
                for depend_script in depend_scripts.split(','):
                    order_list.append(
                        (os.path.basename(script), os.path.basename(depend_script))
                    )

        # write order list
        for order in order_list:
            out.write(write_order(*order))

        # write logdir
        if not os.path.exists(logdir):
            os.makedirs(logdir)

        out.write('\nlog_dir %s\n' % logdir)

        print('make job: {}'.format(outfile))


if __name__ == "__main__":

    main()
